#include "wavAudio.h"

void wavAudio::readInfo(const char* filepath)
{
	if (!is_train && !is_test_accuracy) {
		std::cout << "begin to read info..." << std::endl;
	}
	void* obj = wav_read_open(filepath);

	// data
	int res = wav_get_header(obj, &this->format, &this->channels, &this->sample_rate, &this->bits_per_sample, &this->data_length);
	if (!res)
	{
		std::cerr << "get ref header error: " << res << std::endl;
		return;
	}

	this->samples_num = this->data_length * 8 / bits_per_sample;
	this->data = std::vector<int16_t>(this->samples_num, 0);
	this->process_data = std::vector<int16_t>(this->samples_num, 0);
	res = wav_read_data(obj, reinterpret_cast<unsigned char*>(this->data.data()), this->data_length);
	if (res < 0)
	{
		std::cerr << "read wav file error: " << res << std::endl;
		return;
	}

	std::remove("data/wavInfo.txt");
	std::ofstream outfile("data/wavInfo.txt");
	outfile << "file path:" << filepath << std::endl;
	outfile << "data length:" << this->data_length << std::endl;
	outfile << "audio data format:" << this->format << std::endl;
	outfile << "channels number:" << this->channels << std::endl;
	outfile << "sample rate:" << this->sample_rate << std::endl;
	outfile << "bits per sample:" << this->bits_per_sample << std::endl;
	outfile << "samples number:" << this->samples_num << std::endl << std::endl;

	if (!is_train && !is_test_accuracy) {
		std::cout << "\r" << std::string(80, ' ') << "\r";
	}
	for (int i = 0; i < this->samples_num; i++) {
		if (!is_train && !is_test_accuracy) {
			if (i % 10000 == 0 || i == this->samples_num - 1) {
				displayProgressBar(i + 1, this->samples_num);
			}
		}
		outfile << this->data[i] << " ";
	}
	outfile << std::endl;
	outfile.close();
	wav_read_close(obj);

	if (!is_train && !is_test_accuracy) {
		std::cout << std::endl << "read info done" << std::endl;
	}
}

void wavAudio::processdata()
{
	std::ofstream outfile("data/wavInfo.txt", std::ios::app);

	// 检查样本数是否为零
	if (this->samples_num == 0) {
		std::cerr << "No samples to process." << std::endl;
		return;
	}

	// 初始化最大数据为最小值
	this->max_data = std::numeric_limits<int>::min();

	// 找到最大数据值
	for (int i = 0; i < this->samples_num; i++) {
		int d = this->data[i] > 0 ? this->data[i] : -this->data[i];
		this->max_data = this->max_data > d ? this->max_data : d;
	}
	if (DEBUG) {
		outfile << std::endl << "Max data:" << this->max_data << std::endl;
	}
	// 检查 max_data 是否为零以避免除零错误
	if (this->max_data == 0) {
		std::cerr << "All data points are zero, cannot process data." << std::endl;
		for (int i = 0; i < this->samples_num; i++) {
			this->process_data[i] = 0;
		}
		return;
	}

	this->_per = int(this->max_data * 0.1);

	if (DEBUG) {
		outfile << std::endl << "_per data:" << this->_per << std::endl;
	}
	int index = 0; // 记录第三个非零数据的位置
	bool begin_clock_set = false; // 标记是否已设置 begin_clock

	for (int i = 0; i < this->samples_num; i++) {
		int d = this->data[i] > 0 ? this->data[i] : -this->data[i];
		if (d <= this->_per) {
			this->process_data[i] = 0;
		}
		else {
			this->process_data[i] = int(this->data[i] * PEAK / this->max_data);
			index++;
			if (index == 4 && !begin_clock_set) {
				this->begin_clock = i;
				begin_clock_set = true;
			}
		}
	}

	index = 0; // 记录结束位置
	bool end_clock_set = false; // 标记是否已设置 end_clock

	for (int i = this->samples_num - 1; i >= 0; i--) {
		if (this->process_data[i] != 0) {
			index++;
			if (index == 4 && !end_clock_set) {
				this->end_clock = i;
				end_clock_set = true;
				break;
			}
		}
	}

	if (DEBUG) {
		outfile << std::endl << "Begin pos:" << this->begin_clock << std::endl << "End pos:" << this->end_clock << std::endl << std::endl;

		outfile << "Processed Data:" << std::endl;
		for (int i = 0; i < this->samples_num; i++) {
			outfile << this->process_data[i] << " ";
		}
		outfile << std::endl;
	}
	outfile.close();
}

void wavAudio::calculatePitch()
{
	int index[100] = { 0 };
	for (int i = this->begin_clock; i < this->end_clock; i++) {
		if (this->process_data[i] != 0) {
			index[this->process_data[i] / 100 - 1]++;
		}
	}
	int n = 0;
	int m = INT_MIN;
	for (int i = 0; i < 100; i++) {
		if (index[i] > m) {
			n = i;
			m = index[i];
		}
	}
	this->pitch = n * 100 + 50;
}

void wavAudio::calculateAverageLoudness()
{
	int num = 0;
	int sum = 0;
	for (int i = this->begin_clock; i < this->end_clock; i++) {
		if (this->process_data[i] != 0) {
			num++;
			sum += this->process_data[i] ^ 2;
		}
	}
	if (num != 0) {
		this->average_loudness = int(sqrt(sum / num));
	}
	else {
		this->average_loudness = 0;
	}
}

void wavAudio::calculateTime()
{
	this->duration = float(this->end_clock - this->begin_clock) / float(this->sample_rate * this->channels);
}

void wavAudio::calculateEnergy()
{
	double energy = 0;
	float d = 1 / float(this->sample_rate * this->channels);
	for (int i = this->begin_clock; i < this->end_clock; i++) {
		if (this->process_data[i] != 0) {
			energy += (this->process_data[i] ^ 2) * d;
		}
	}
	this->energy = energy;
}

void wavAudio::calculateSTFT()
{
	std::vector<float> x(this->samples_num);
	std::transform(this->process_data.begin(), this->process_data.end(), x.begin(),
		[](int16_t a) {
			return static_cast<float>(a) / 32767.f;
		});
	int n_fft = 400;
	int n_hop = 160;
	int n_mel = 40;
	int fmin = 80;
	int fmax = 7600;

	//std::vector<std::vector<std::complex<float>>> X = librosa::Feature::stft(x, n_fft, n_hop, "hann", true, "reflect");

	std::vector<float> x1;
	for (int i = 0, index = 0; i < DATA_NUM && index < x.size(); index++) {
		if (x[index] != 0) {
			x1.push_back(x[index]);
			i++;
		}
	}
	std::vector<std::vector<Complex>> X = simplifiedSTFT(x1, n_fft, n_hop, "hann", true, "reflect");
	
	// 将实部和虚部分别存储在不同的数组中
	std::vector<float> real_part;
	std::vector<float> imag_part;
	for (const auto& row : X) {
		for (const auto& val : row) {
			real_part.push_back(val.real());
			imag_part.push_back(val.imag());
		}
	}

	// 计算实部的平均值、方差、最大值和最小值
	this->sftf_real_mean = std::accumulate(real_part.begin(), real_part.end(), 0.0f) / real_part.size();
	float sum_of_squares_real = std::inner_product(real_part.begin(), real_part.end(), real_part.begin(), 0.0f);
	this->sftf_real_variance = sum_of_squares_real / real_part.size();
	this->sftf_real_min = *std::min_element(real_part.begin(), real_part.end());
	this->sftf_real_max = *std::max_element(real_part.begin(), real_part.end());

	// 计算虚部的平均值、方差、最大值和最小值
	this->sftf_imag_mean = std::accumulate(imag_part.begin(), imag_part.end(), 0.0f) / imag_part.size();
	float sum_of_squares_imag = std::inner_product(imag_part.begin(), imag_part.end(), imag_part.begin(), 0.0f);
	this->sftf_imag_variance = sum_of_squares_imag / imag_part.size();
	this->sftf_imag_min = *std::min_element(imag_part.begin(), imag_part.end());
	this->sftf_imag_max = *std::max_element(imag_part.begin(), imag_part.end());

	this->sftfRealArray = std::vector<float>(STFT_NUM, 0);
	this->sftfImagArray = std::vector<float>(STFT_NUM, 0);
	int index = 0;
	for (int i = 0; i < STFT_NUM; i++) {
		while (real_part[index] == 0 && imag_part[index] == 0) {
			index++;
			if (index >= real_part.size() || index >= imag_part.size()) {
				index--;
				break;
			}
		}
		if (index < real_part.size() && index < imag_part.size()) {
			this->sftfRealArray[i] = real_part[index];
			this->sftfImagArray[i] = imag_part[index];
		}
	}
}

void wavAudio::calculateMelspectrogram()
{
	std::vector<float> x(this->samples_num);
	std::transform(this->process_data.begin(), this->process_data.end(), x.begin(),
		[](int16_t a) {
			return static_cast<float>(a) / 32767.f;
		});

	// 计算梅尔频谱图
	int n_fft = 400;
	int n_hop = 160;
	int n_mel = 40;
	int fmin = 80;
	int fmax = 7600;

	//std::vector<std::vector<float>> mels = librosa::Feature::melspectrogram(x, this->sample_rate, n_fft, n_hop, "hann", true, "reflect", 2.f, n_mel, fmin, fmax);

	std::vector<float> x1;
	for (int i = 0, index = 0; i < DATA_NUM && index < x.size(); index++) {
		if (x[index] != 0) {
			x1.push_back(x[index]);
			i++;
		}
	}
	std::vector<std::vector<float>> mels = melSpectrogram(x1, sample_rate, n_fft, n_hop, "hann", true, "reflect", 2.f, n_mel, fmin, fmax);

	// 将二维数组打平为一维数组
	std::vector<float> flattenedMels;
	for (const auto& row : mels) {
		flattenedMels.insert(flattenedMels.end(), row.begin(), row.end());
	}

	// 计算平均值
	this->mels_mean = std::accumulate(flattenedMels.begin(), flattenedMels.end(), 0.0f) / flattenedMels.size();

	// 计算方差
	float sum_of_squares = 0.0f;
	for (float value : flattenedMels) {
		sum_of_squares += (value - this->mels_mean) * (value - this->mels_mean);
	}
	this->mels_variance = sum_of_squares / flattenedMels.size();
	// 计算最小值和最大值
	this->mels_min = *std::min_element(flattenedMels.begin(), flattenedMels.end());
	this->mels_max = *std::max_element(flattenedMels.begin(), flattenedMels.end());

	// calculate the top MELS_NUM dimension
	this->melsArray = std::vector<float>(MELS_NUM, 0);
	int index = 0;
	for (int i = 0; i < MELS_NUM; i++) {
		while (flattenedMels[index] == 0)
		{
			index++;
			if (index >= flattenedMels.size()) {
				index--;
				break;
			}
		}
		if (index < flattenedMels.size()) {
			this->melsArray[i] = flattenedMels[index];
		}
	}
}


void wavAudio::computeCharacter()
{
	if (!is_train && !is_test_accuracy) {
		std::cout << "begin to calculate character..." << std::endl;
	}
	this->processdata();
	if (!is_train && !is_test_accuracy) {
		std::cout << "\r" << std::string(80, ' ') << "\r";
		displayProgressBar(1, 7);
	}
	this->calculatePitch();
	if (!is_train && !is_test_accuracy) {
		displayProgressBar(2, 7);
	}
	this->calculateAverageLoudness();
	if (!is_train && !is_test_accuracy) {
		displayProgressBar(3, 7);
	}
	this->calculateTime();
	if (!is_train && !is_test_accuracy) {
		displayProgressBar(4, 7);
	}
	this->calculateEnergy();
	if (!is_train && !is_test_accuracy) {
		displayProgressBar(5, 7);
	}
	this->calculateSTFT();
	if (!is_train && !is_test_accuracy) {
		displayProgressBar(6, 7);
	}
	this->calculateMelspectrogram();
	if (!is_train && !is_test_accuracy) {
		displayProgressBar(7, 7);
		std::cout << std::endl << "calculate character done" << std::endl;
	}
}

std::vector<float> wavAudio::returnCharacter()
{
	std::vector<float> res;

	res.push_back(this->duration);
	res.push_back(this->pitch);
	res.push_back(this->average_loudness);
	res.push_back(this->energy);
	res.push_back(this->mels_mean);
	res.push_back(this->mels_variance);
	res.push_back(this->mels_min);
	res.push_back(this->mels_max);
	res.push_back(this->sftf_real_mean);
	res.push_back(this->sftf_real_variance);
	res.push_back(this->sftf_real_max);
	res.push_back(this->sftf_real_min);
	res.push_back(this->sftf_imag_mean);
	res.push_back(this->sftf_imag_variance);
	res.push_back(this->sftf_imag_max);
	res.push_back(this->sftf_imag_min);
	for (int i = 0; i < MELS_NUM; i++) {
		res.push_back(this->melsArray[i]);
	}
	for (int i = 0; i < STFT_NUM; i++) {
		res.push_back(this->sftfRealArray[i]);
		res.push_back(this->sftfImagArray[i]);
	}

	return res;
}

wavAudio::~wavAudio()
{
}
